"""
https://blog.csdn.net/menghaocheng/article/details/102783705

【TF2.0-CNN】迁移学习（将inceptionV3应用到猫狗分类）
"""

from tensorflow.keras.preprocessing.image import ImageDataGenerator
import tensorflow as tf
import numpy as np
import os


def sep(label = '', cnt=32):
    print('-' * cnt, label, '-' * cnt, sep='')


sep('Config')
tf.random.set_seed(1)
np.random.seed(1)

VER = 'v4.0'
FILE_NAME = os.path.basename(__file__)
LOG_DIR = os.path.join('_log', FILE_NAME, VER)
if os.path.exists(LOG_DIR) and len(os.listdir(LOG_DIR)) > 0:
    NEED_LOG = False
else:
    NEED_LOG = True
print('VER:', VER)
print('NEED_LOG:', NEED_LOG)

sep('Load model')
INPUT_SHAPE = (150, 150, 3)
local_weights_file = '../../../../large_data/model/inceptionV3/inception_v3_weights_tf_dim_ordering_tf_kernels_notop.h5'

pre_trained_model = tf.keras.applications.inception_v3.InceptionV3(input_shape=INPUT_SHAPE,
                                                                   include_top=False,
                                                                   weights=None)

pre_trained_model.load_weights(local_weights_file)
for layer in pre_trained_model.layers:
    layer.trainable = False

last_layer = pre_trained_model.get_layer('mixed7')
print('last layer output shape: ', last_layer.output_shape)
last_output = last_layer.output

pre_trained_model.summary()

sep('For graph')
if NEED_LOG:
    fw = tf.summary.create_file_writer(LOG_DIR)
    with fw.as_default():
        tf.summary.trace_on()

imgs = np.zeros([1, *INPUT_SHAPE], dtype=np.float32)
vecs = pre_trained_model(imgs)
print('vecs', vecs.shape)

if NEED_LOG:
    with fw.as_default():
        tf.summary.trace_export('graph', 0)
    fw.close()

